#include "defines.h"
#include "net.h"
#include "timer.h"
#include "lib.h"
#include "nic.h"

static void show_progress(uint16 blocknum)
{
    if ((blocknum & 0x7) == 1)
        puts("*");
}

static void tftp_rrq(uint8 *buffer, net_params *np, const char *filename)
{
    uint16 len;
    uint16 *p = (uint16 *)(buffer + sizeof(ethernet_header) + sizeof(ip_header) + sizeof(udp_header));

    *p++ = 1;
    len = strlen(filename) + 1;
    memcpy(p, filename, len);
    strcpy((char *)p + len, "octet");

    send_by_udp(buffer, 2 + len + 6, np);
}

static void tftp_ack(uint8 *buffer, net_params *np, uint16 blocknum)
{
    uint16 *p = (uint16 *)(buffer + sizeof(ethernet_header) + sizeof(ip_header) + sizeof(udp_header));

    *p++ = 4;
    *p = blocknum;

    send_by_udp(buffer, 4, np);
}

static void tftp_error(uint8 *buffer, net_params *np, uint16 code, const char *msg)
{
    uint16 len;
    uint16 *p = (uint16 *)(buffer + sizeof(ethernet_header) + sizeof(ip_header) + sizeof(udp_header));

    *p++ = 5;
    *p++ = code;
    len = strlen(msg) + 1;
    memcpy(p, msg, len);
    send_by_udp(buffer, 4 + len, np);
}

#define RETRY           4
#define RETRY_INTERVAL  2

int tftp_down(uint8 *buffer, net_params *np, const char *filename,
              uint8 *downbuf, uint32 downsize, uint32 *sizep)
{
    ethernet_header *ehdr = (ethernet_header *)buffer;
    arp_packet *apacket = (arp_packet *)(buffer + sizeof(*ehdr));
    ip_header *ihdr = (ip_header *)(buffer + sizeof(*ehdr));

    uint16 block_number = 1;
    int retry, ret = 0;
    uint32 size = 0;

    timeout_init(RETRY_INTERVAL);

    np->dst_port = 69;
    retry = 0;
    tftp_rrq(buffer, np, filename);
    timeout_start();

    // ブロック1の受信
    while (1) {
        if (is_timeout()) {
            if (++retry > RETRY) {
                puts("TFTP timeout.");
                ret = -1;
                goto end;
            }
            tftp_rrq(buffer, np, filename);
            timeout_start();
        }

        if (frame_receive(buffer) == 0)
        {
            if ((ehdr->eth_ethernet_type == 0x0806)
                    && !memcmp(apacket->arp_dst_IP, np->src_IP, 4)
                    && (apacket->arp_operation == 1))
            {

                arp_reply(buffer, np);

            } else if ((ehdr->eth_ethernet_type == 0x0800)
                    && !memcmp(ihdr->ip_dst_IP, np->src_IP, 4)
                    && (ihdr->ip_protocol == 17))
            {
                udp_header *uhdr = (udp_header *)((uint8 *)ihdr + ((ihdr->ip_version_length & 0x0f) << 2));
                uint8 *data = (uint8 *)uhdr + sizeof(*uhdr);

                if (uhdr->udp_dst_port == np->src_port)
                {
                    if (ip_checksum((uint8 *)ihdr, (ihdr->ip_version_length & 0x0f) << 2) != 0) {
                        continue;
                    }
                    if (uhdr->udp_checksum) { // チェックサム不使用でも特に何もしないことに注意
                        if (udp_checksum((uint8 *)uhdr, uhdr->udp_length, ihdr->ip_src_IP, ihdr->ip_dst_IP) != 0xffff) {
                            continue;
                        }
                    }

                    uint16 opcode = *(uint16 *)data;
                    if (opcode == 3 && *(uint16 *)(data + 2) == block_number) {
                        uint16 dsize;

                        np->dst_port = uhdr->udp_src_port;
                        dsize = uhdr->udp_length - sizeof(*uhdr) - 4;

                        if (downsize < size + dsize) {
                            dsize = downsize - size;
                            memcpy(downbuf + size, data + 4, dsize);
                            size += dsize;

                            puts("Download buffer overflow.");
                            tftp_error(buffer, np, 3, "disk is full");

                            ret = -1;
                            goto end;
                        }

                        tftp_ack(buffer, np, block_number++);
                        memcpy(downbuf + size, data + 4, dsize);
                        size += dsize;
                        show_progress(block_number - 1);

                        if (dsize < 512) {
                            goto end;
                        } else {
                            retry = 0;
                            timeout_start();
                            break;
                        }

                    } else if (opcode == 5) {
                        puts("Error received (");
                        putxval(*(uint16 *)(data + 2), 0);
                        puts(", ");
                        puts(data + 4);
                        puts(")");
                        ret = -1;
                        goto end;
                    }
                }
            }
        }
    }

    // ブロック2以降の受信
    while (1) {
        if (is_timeout()) {
            if (++retry > RETRY) {
                tftp_error(buffer, np, 0, "timeout");
                puts("TFTP timeout.");
                ret = -1;
                goto end;
            }
            tftp_ack(buffer, np, block_number - 1);
            timeout_start();
        }

        if (frame_receive(buffer) == 0) {
            if ((ehdr->eth_ethernet_type == 0x0806)
                    && !memcmp(apacket->arp_dst_IP, np->src_IP, 4)
                    && (apacket->arp_operation == 1))
            {

                arp_reply(buffer, np);

            } else if ((ehdr->eth_ethernet_type == 0x0800)
                    && !memcmp(ihdr->ip_dst_IP, np->src_IP, 4)
                    && (ihdr->ip_protocol == 17))
            {
                udp_header *uhdr = (udp_header *)((uint8 *)ihdr + ((ihdr->ip_version_length & 0x0f) << 2));
                uint8 *data = (uint8 *)uhdr + sizeof(*uhdr);

                if (uhdr->udp_dst_port == np->src_port && uhdr->udp_src_port == np->dst_port)
                {
                    if (ip_checksum((uint8 *)ihdr, (ihdr->ip_version_length & 0x0f) << 2) != 0) {
                        continue;
                    }
                    if (uhdr->udp_checksum) {
                        if (udp_checksum((uint8 *)uhdr, uhdr->udp_length, ihdr->ip_src_IP, ihdr->ip_dst_IP) != 0xffff) {
                            continue;
                        }
                    }

                    uint16 opcode = *(uint16 *)data;
                    if (opcode == 3) {

                        if (*(uint16 *)(data + 2) != block_number) {
                            puts("Bad block number.");
                            tftp_error(buffer, np, 0, "bad block number");
                            ret = -1;
                            goto end;
                        }

                        uint16 dsize = uhdr->udp_length - sizeof(*uhdr) - 4;

                        if (downsize < size + dsize) {
                            dsize = downsize - size;
                            memcpy(downbuf + size, data + 4, dsize);
                            size += dsize;

                            puts("Download buffer overflow.");
                            tftp_error(buffer, np, 3, "disk is full");

                            ret = -1;
                            goto end;
                        }

                        tftp_ack(buffer, np, block_number++);
                        memcpy(downbuf + size, data + 4, dsize);
                        size += dsize;
                        show_progress(block_number - 1);

                        if (dsize < 512) {
                            goto end;
                        } else {
                            retry = 0;
                            timeout_start();
                        }

                    } else if (opcode == 5) {
                        puts("Error received (");
                        putxval(*(uint16 *)(data + 2), 0);
                        puts(", ");
                        puts(data + 4);
                        puts(")");
                        ret = -1;
                        goto end;

                    } else {
                        puts("Bad operation.");
                        tftp_error(buffer, np, 4, "bad operation");
                        ret = -1;
                        goto end;
                    }
                }
            }
        }
    }

end:
    puts("\n");
    timeout_end();
    *sizep = size;
    return ret;
}
